import logging
import queue
import socket
import sys
import threading
import time
from collections import namedtuple
from typing import Optional

import numpy
import select as sel

# --- 全局辅助定义 (这些不依赖于客户端实例) ---
Point = namedtuple("Point", "x y")
Rect = namedtuple("Rect", "x y width height")
RecvNullExc = type("RecvNullExc", (Exception,), {})

# 修正像素格式定义，使用BGRX和BGRA格式确保原色显示
pixel_formats = {
    # BGRX格式：32位/像素，24位深度，无Alpha通道
    "bgrx": b"\x20\x18\x00\x01\x00\xff\x00\xff\x00\xff\x00\x08\x10\x00\x00\x00",
    # BGRA格式：32位/像素，32位深度，包含Alpha通道
    "bgra": b"\x20\x20\x00\x01\x00\xff\x00\xff\x00\xff\x18\x08\x10\x00\x00\x00",
}
encodings = {0}  # 只使用 Raw 编码


def _read(sock: socket.socket, length: int) -> bytes:
    """从 socket 读取指定长度的数据，增加超时处理。"""
    data = b""
    end_time = time.time() + 5  # 5秒超时
    while len(data) < length:
        if time.time() > end_time:
            raise socket.timeout("Read operation timed out")
        try:
            chunk = sock.recv(length - len(data))
            if not chunk:
                raise RecvNullExc("Socket connection closed by peer while reading")
            data += chunk
        except BlockingIOError:
            # 在非阻塞模式下可能发生，稍等重试
            time.sleep(0.01)
        except ConnectionResetError:
            raise
    return data


def _read_int(sock: socket.socket, length: int) -> int:
    """从 socket 读取一个整数"""
    return int.from_bytes(_read(sock, length), "big")


class VNCClient:
    """
    一个VNC客户端，内部使用命令队列和专用网络线程来处理通信，
    实现了高效、无阻塞的屏幕捕获和远程控制，支持原色显示。
    """

    def __init__(self, ip="127.0.0.1", port=8801, timeout=5, capture_rect: Optional[Rect] = None,
                 password: Optional[str] = None, use_alpha=False):
        """
        初始化VNC客户端。
        :param ip: VNC服务器IP地址
        :param port: VNC服务器端口
        :param timeout: 连接和读取的超时时间
        :param capture_rect: (可选) 一个Rect对象(x, y, width, height)，定义了始终捕获的区域。
                             如果为None，则捕获全屏。
        :param password: (可选) VNC服务器密码，如果为None则使用无密码连接
        :param use_alpha: 是否使用包含Alpha通道的格式(BGRA)，默认不使用(BGRX)
        """
        self.ip = ip
        self.port = port
        self.timeout = timeout
        self.password = password
        self.use_alpha = use_alpha  # 新增参数控制是否使用Alpha通道

        # 正确地存储用户请求的捕获区域
        self._requested_capture_rect = capture_rect

        # 将这些属性的最终赋值推迟到connect方法中
        self.full_screen_rect: Optional[Rect] = None
        self.capture_rect: Optional[Rect] = None

        # --- 核心数据和线程组件 ---
        self.sock: Optional[socket.socket] = None
        self.pixels: numpy.ndarray | None = None  # 图像数据缓冲区
        self._pixels_lock = threading.Lock()  # 保护self.pixels的锁

        self._command_queue = queue.Queue()  # 线程安全的命令队列
        self._signal_sock_r, self._signal_sock_w = socket.socketpair()

        self.stop_event = threading.Event()  # 全局停止信号
        self._network_thread: Optional[threading.Thread] = None  # 网络线程句柄

        # 标志位，用于确认是否已收到第一帧有效数据
        self._initial_frame_received = False

    def is_ready(self) -> bool:
        """检查VNC客户端是否已接收到第一帧有效数据并准备好提供图像。"""
        return self._initial_frame_received

    def connect(self):
        """连接到VNC服务器，完成握手，并启动核心网络线程。"""
        logging.info(f"正在连接到 {self.ip}:{self.port}...")
        self.sock = socket.create_connection((self.ip, self.port), timeout=self.timeout)
        self.sock.setblocking(False)  # 将socket设为非阻塞模式，以配合select
        logging.info("连接已建立。")

        self._perform_handshake()  # 这个方法会设置 self.full_screen_rect

        # 在获取到全屏尺寸后，决定最终的捕获区域
        if self._requested_capture_rect is None:
            # 如果用户没有指定区域，就使用服务器返回的全屏尺寸
            self.capture_rect = self.full_screen_rect
        else:
            # 否则，使用用户指定的区域
            self.capture_rect = self._requested_capture_rect

        # 添加一个安全检查
        if self.capture_rect is None:
            raise RuntimeError("无法确定捕获区域。服务器可能未返回屏幕尺寸。")

        logging.info(f"最终捕获区域: {self.capture_rect}")

        # 根据是否使用Alpha通道初始化不同的缓冲区
        channels = 4 if self.use_alpha else 3
        self.pixels = numpy.zeros((self.capture_rect.height, self.capture_rect.width, channels), "B")

        # 启动网络核心循环线程
        logging.info("握手完成，启动网络核心线程...")
        self.stop_event.clear()
        self._network_thread = threading.Thread(target=self._network_loop, daemon=True)
        self._network_thread.start()

    def _network_loop(self):
        """网络核心循环，在专用线程中运行，处理socket通信和命令队列。"""
        try:
            # 连接后，立即请求一次全量更新来获取初始画面
            self._send_frame_request(incremental=False)

            while not self.stop_event.is_set():
                # 使用select监听socket的可读事件和信号socket的可读事件
                read_list = [self.sock, self._signal_sock_r]
                readable, _, _ = sel.select(read_list, [], [], 0.016)

                for s in readable:
                    if s is self.sock:
                        # 如果是socket可读，说明服务器发来了帧更新数据
                        self._process_framebuffer_update()
                    elif s is self._signal_sock_r:
                        # 从信号socket读取，清空信号
                        self._signal_sock_r.recv(1)
                        self._process_command_queue()

                # 如果select超时，发送一次增量请求作为心跳
                if not readable:
                    self._send_frame_request(incremental=True)

        except (ConnectionError, RecvNullExc, BrokenPipeError, socket.timeout) as e:
            logging.warning(f"VNC网络线程因错误退出: {e}")
        finally:
            logging.info("VNC网络核心线程结束。")
            # 确保关闭信号socket
            self._signal_sock_r.close()
            self._signal_sock_w.close()

            self.stop_event.set()  # 通知所有其他部分也停止

    def _process_command_queue(self):
        """处理命令队列中的所有待发命令。"""
        while not self._command_queue.empty():
            try:
                command = self._command_queue.get_nowait()
                method_name, args, kwargs = command
                # 动态调用对应的内部发送方法
                method = getattr(self, f"_send_{method_name}")
                method(*args, **kwargs)
            except queue.Empty:
                break
            except Exception as e:
                logging.warning(f"处理命令时出错: {e}")

    def _wake_up_network_thread(self):
        """向信号socket写入一个字节，以唤醒阻塞的 select()"""
        if self._signal_sock_w and not self.stop_event.is_set():
            try:
                self._signal_sock_w.send(b'x')
            except (OSError, BrokenPipeError):
                # 如果socket已关闭，则忽略
                pass

    # --- 以下是供外部调用的公共API ---

    def key_down(self, key: str):
        """将'按下按键'命令放入队列"""
        self._command_queue.put(('key_event', (key, True), {}))
        self._wake_up_network_thread()

    def key_up(self, key: str):
        """将'松开按键'命令放入队列"""
        self._command_queue.put(('key_event', (key, False), {}))
        self._wake_up_network_thread()

    def move(self, x: int, y: int):
        """将'鼠标移动'命令放入队列"""
        self._command_queue.put(('mouse_event', (Point(x, y), 0), {}))
        self._wake_up_network_thread()

    def click(self, x: int, y: int):
        """将'鼠标点击'(按下+松开)命令放入队列"""
        pos = Point(x, y)
        self._command_queue.put(('mouse_event', (pos, 1), {}))  # 按下左键
        self._command_queue.put(('mouse_event', (pos, 0), {}))  # 松开所有键
        self._wake_up_network_thread()

    def mouse_down(self, x: int, y: int, button_mask: int = 1):
        """将'鼠标按下'命令放入队列。"""
        pos = Point(x, y)
        self._command_queue.put(('mouse_event', (pos, button_mask), {}))
        self._wake_up_network_thread()

    def mouse_up(self, x: int, y: int):
        """将'鼠标松开'命令放入队列。"""
        pos = Point(x, y)
        # 松开所有按钮
        self._command_queue.put(('mouse_event', (pos, 0), {}))
        self._wake_up_network_thread()

    def scroll_wheel(self, x: int, y: int, direction: int):
        """将'鼠标滚轮'命令放入队列。"""
        pos = Point(x, y)
        # 滚轮事件使用正确的按钮掩码
        # 向上滚动: 0x08, 向下滚动: 0x10
        button_mask = 0x08 if direction > 0 else 0x10

        # 滚轮事件需要先按下再松开
        self._command_queue.put(('mouse_event', (pos, button_mask), {}))
        self._wake_up_network_thread()
        # 短暂延迟后松开
        self._command_queue.put(('mouse_event', (pos, 0), {}))
        self._wake_up_network_thread()

    def scroll_horizontal(self, x: int, y: int, direction: int):
        """将'鼠标水平滚轮'命令放入队列。"""
        pos = Point(x, y)
        # 水平滚轮事件
        # 向右滚动: 0x20, 向左滚动: 0x40
        button_mask = 0x20 if direction > 0 else 0x40

        # 滚轮事件需要先按下再松开
        self._command_queue.put(('mouse_event', (pos, button_mask), {}))
        self._wake_up_network_thread()
        # 短暂延迟后松开
        self._command_queue.put(('mouse_event', (pos, 0), {}))
        self._wake_up_network_thread()

    def get_frame(self, copy: bool = True) -> numpy.ndarray:
        """
        线程安全地从缓冲区获取最新的帧图像。
        默认返回一个副本以避免多线程问题。
        """
        if self.pixels is None:
            raise ConnectionError("VNC未连接或尚未收到第一帧图像。")
        with self._pixels_lock:
            return self.pixels.copy() if copy else self.pixels  # 根据参数决定是否复制

    # --- 以下是网络线程专用的私有方法 ---

    def _send_frame_request(self, incremental: bool):
        """发送帧更新请求到socket (只应被网络线程调用)"""
        self.sock.sendall(
            b"\x03" + (b"\x01" if incremental else b"\x00") +
            self.capture_rect.x.to_bytes(2, 'big') +
            self.capture_rect.y.to_bytes(2, 'big') +
            self.capture_rect.width.to_bytes(2, 'big') +
            self.capture_rect.height.to_bytes(2, 'big')
        )

    def _process_framebuffer_update(self):
        msg_type = _read_int(self.sock, 1)
        if msg_type != 0: return

        _read(self.sock, 1)
        num_rects = _read_int(self.sock, 2)

        for _ in range(num_rects):
            area_rect_global = Rect(
                _read_int(self.sock, 2), _read_int(self.sock, 2),
                _read_int(self.sock, 2), _read_int(self.sock, 2)
            )
            area_encoding = _read_int(self.sock, 4)
            if area_encoding != 0: continue

            # 读取像素数据并整形
            chunk_size = area_rect_global.height * area_rect_global.width * 4
            area_data = _read(self.sock, chunk_size)
            area_array = numpy.frombuffer(area_data, dtype="B").reshape(
                (area_rect_global.height, area_rect_global.width, 4)
            )

            # 关键修正：将RGBA转换为BGRA或BGR
            # 交换R和B通道以获取正确的原色
            if self.use_alpha:
                # 保留Alpha通道，转换为BGRA
                area_array = area_array[..., [2, 1, 0, 3]]  # B, G, R, A
            else:
                # 不使用Alpha通道，转换为BGR
                area_array = area_array[..., [2, 1, 0]]  # B, G, R

            # 计算服务器更新区域和我们关心的捕获区域之间的交集
            ix = max(self.capture_rect.x, area_rect_global.x)
            iy = max(self.capture_rect.y, area_rect_global.y)
            ir = min(self.capture_rect.x + self.capture_rect.width, area_rect_global.x + area_rect_global.width)
            ib = min(self.capture_rect.y + self.capture_rect.height, area_rect_global.y + area_rect_global.height)

            iw = ir - ix
            ih = ib - iy

            if iw > 0 and ih > 0:
                # 计算数据源(area_array)中需要复制的区域
                src_x = ix - area_rect_global.x
                src_y = iy - area_rect_global.y
                src_slice = (slice(src_y, src_y + ih), slice(src_x, src_x + iw))

                # 计算目标(self.pixels)中需要粘贴的区域
                dest_x = ix - self.capture_rect.x
                dest_y = iy - self.capture_rect.y
                dest_slice = (slice(dest_y, dest_y + ih), slice(dest_x, dest_x + iw))

                # 执行复制，更新共享的图像缓冲区时必须加锁
                with self._pixels_lock:
                    self.pixels[dest_slice] = area_array[src_slice]

                    # 设置标志位，确认已收到第一帧有效数据
                    if not self._initial_frame_received:
                        self._initial_frame_received = True
                        logging.info("网络线程：已接收并处理第一帧有效像素数据（原色）")

    def _perform_handshake(self):
        """执行VNC握手过程，包括版本协商、认证和像素格式设置"""
        self.sock.setblocking(True)  # 临时改回阻塞模式以便于握手
        try:
            intro = self.sock.recv(12)
            assert intro[:4] == b"RFB "
            self.sock.sendall(b"RFB 003.008\n")

            # 读取认证类型
            num_auth_types = _read_int(self.sock, 1)
            auth_types = _read(self.sock, num_auth_types)  # read auth types

            # 检查是否有密码认证类型 (2) 和无密码认证类型 (1)
            has_password_auth = b"\x02" in auth_types
            has_none_auth = b"\x01" in auth_types

            if self.password and has_password_auth:
                # 使用密码认证
                self.sock.sendall(b"\x02")  # choose "VNC Authentication" auth type

                # 读取16字节的挑战
                challenge = _read(self.sock, 16)
                logging.info(f"收到VNC认证挑战: {challenge.hex()}")

                # 使用DES加密密码
                response = self._encrypt_challenge(challenge, self.password)
                self.sock.sendall(response)

                # 读取认证结果
                auth_result = _read_int(self.sock, 4)
                if auth_result == 0:  # 认证成功
                    logging.info("VNC密码认证成功")
                else:
                    # 读取错误信息
                    if auth_result == 1:
                        logging.error("VNC密码认证失败：密码错误")
                    elif auth_result == 2:
                        logging.error("VNC密码认证失败：服务器拒绝连接")
                    else:
                        logging.error(f"VNC密码认证失败：未知错误代码 {auth_result}")
                    raise ConnectionError("VNC密码认证失败")
            elif not self.password and has_none_auth:
                # 使用无密码认证
                self.sock.sendall(b"\x01")  # choose "None" auth type
                auth_result = _read_int(self.sock, 4)
                if auth_result == 0:  # 认证成功
                    logging.info("VNC无密码认证成功")
                else:
                    raise ConnectionError("VNC无密码认证失败")
            else:
                # 如果服务器不支持请求的认证类型
                if self.password and not has_password_auth:
                    raise ConnectionError("VNC服务器不支持密码认证，但提供了密码")
                elif not self.password and not has_none_auth:
                    raise ConnectionError("VNC服务器不支持无密码认证，但未提供密码")
                else:
                    raise ConnectionError("VNC认证类型不匹配")

            self.sock.sendall(b"\x01")  # Share desktop flag
            width, height = _read_int(self.sock, 2), _read_int(self.sock, 2)
            self.full_screen_rect = Rect(0, 0, width, height)
            _read(self.sock, 16)  # read pixel format
            _read(self.sock, _read_int(self.sock, 4))  # read name string

            # 设置我们期望的像素格式 - 使用BGRX或BGRA确保原色显示
            format_key = "bgra" if self.use_alpha else "bgrx"
            desired_format = pixel_formats[format_key]
            self.sock.sendall(b"\x00\x00\x00\x00" + desired_format)
            logging.info(f"请求 {format_key.upper()} 像素格式（确保原色显示）")

            # 设置编码方式
            enc_list = sorted(list(encodings))
            self.sock.sendall(
                b"\x02\x00" + len(enc_list).to_bytes(2, "big") +
                b"".join(enc.to_bytes(4, "big", signed=True) for enc in enc_list)
            )
        finally:
            self.sock.setblocking(False)

    def _encrypt_challenge(self, challenge, password):
        """使用DES加密VNC认证挑战"""
        try:
            from Crypto.Cipher import DES
        except ImportError:
            # 如果没有pycryptodome，尝试使用内置的简单实现
            return self._simple_vnc_encrypt(challenge, password)

        # VNC使用DES加密，但密钥有特殊处理
        # 密码需要截断或填充到8字节，并且位需要反转
        key = (password + '\x00' * 8)[:8].encode('latin1')

        # VNC特殊要求：反转每个字节的位
        key = bytes([self._reverse_bits(b) for b in key])

        # 使用DES ECB模式加密
        cipher = DES.new(key, DES.MODE_ECB)
        response = cipher.encrypt(challenge)

        logging.info(f"DES加密响应: {response.hex()}")
        return response

    def _reverse_bits(self, byte_val):
        """反转字节的位"""
        result = 0
        for i in range(8):
            if byte_val & (1 << i):
                result |= (1 << (7 - i))
        return result

    def _simple_vnc_encrypt(self, challenge, password):
        """简单的VNC加密实现（如果没有Crypto库）"""
        # 这是一个简化的实现，实际VNC需要DES加密
        # 在生产环境中应该安装pycryptodome
        logging.warning("使用简化的VNC加密，建议安装pycryptodome库")

        # 简单的XOR加密作为后备方案
        key = (password + '\x00' * 8)[:8].encode('latin1')
        response = bytearray(16)

        for i in range(16):
            response[i] = challenge[i] ^ key[i % 8]

        return bytes(response)

    def close(self):
        """优雅地关闭VNC客户端和所有相关资源。"""
        if self.stop_event.is_set():
            return  # 防止重复关闭

        logging.info("正在关闭VNC连接...")
        self.stop_event.set()

        # 唤醒可能在select中等待的网络线程
        self._wake_up_network_thread()

        if self._network_thread and self._network_thread.is_alive():
            self._network_thread.join(timeout=2)  # 等待网络线程结束

        if self.sock:
            self.sock.close()
            self.sock = None

        # 关闭信号socket
        try:
            self._signal_sock_w.close()
        except (OSError, AttributeError):
            pass
        try:
            self._signal_sock_r.close()
        except (OSError, AttributeError):
            pass

        logging.info("VNC连接已关闭。")


# 示例使用代码
def main():
    import cv2

    # 配置日志
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    # 创建VNC客户端实例，默认使用BGR格式(无Alpha)
    # 如需使用Alpha通道，添加 use_alpha=True 参数
    client = VNCClient(
        ip="127.0.0.1",  # 替换为你的VNC服务器IP
        port=5900,  # 替换为你的VNC服务器端口
        password=None  # 替换为你的VNC密码，无密码则为None
    )

    try:
        # 连接到VNC服务器
        client.connect()

        # 等待接收第一帧数据
        logging.info("等待接收第一帧图像...")
        while not client.is_ready():
            time.sleep(0.1)

        logging.info("开始显示VNC画面，按ESC键退出")
        # 显示VNC画面
        while True:
            # 获取当前帧
            frame = client.get_frame()

            # 显示图像
            cv2.imshow("VNC 原色画面", frame)

            # 按ESC键退出
            key = cv2.waitKey(1)
            if key == 27:  # ESC键
                break

    except Exception as e:
        logging.error(f"发生错误: {e}")
    finally:
        # 关闭资源
        client.close()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
